#include <iostream>
#include <limits>
#include <thread>
#include "VkBundle.h"
#include <GLFW/glfw3.h>
#include "VkHelper.h"
#include "Geometry.h"
#include "GeometryDemo.h"
#include "Texture.h"

VkBundle vk;    //vulkan context
#define MAX_FRAMES_COUNT  2 //双缓冲

const char *defaultShader = "demo01";
CGeometry geometry;
CTexture texture1;
CTexture texture2;

#define WINDOW_WIDTH 800
#define WINDOW_HEIGHT 600
GLFWwindow *window;
UniformObject ubo{};

void createWindow();
void init();
void initImages();
void initGeometry();
void renderLoop();
void updateUniformBuffer(uint32_t currentFrame);
void updateDescriptorSets(uint32_t currentFrame);
void reCreateSwapchain();
void destroy();

int main() {
    std::cout << "Hello, Learn Vulkan!" << std::endl;

    init();
    renderLoop();
    destroy();

    return 0;
}


void createSwapchainAndResource(){
    VkHelper::createSwapchain(vk.deviceInfo.physicalDev, vk.deviceInfo.device, vk.surface, vk.queueInfo.workQueueIndex, vk.queueInfo.presentQueueIndex,
                              &vk.swapchainParam,&vk.swapchain);
    CALL_VK(vkGetSwapchainImagesKHR(vk.deviceInfo.device, vk.swapchain, &vk.swapchainImage.imageCount, nullptr));
    vk.swapchainImage.images = static_cast<VkImage *>(malloc(sizeof(VkImage) * vk.swapchainImage.imageCount));
    CALL_VK(vkGetSwapchainImagesKHR(vk.deviceInfo.device, vk.swapchain, &vk.swapchainImage.imageCount, vk.swapchainImage.images));
    vk.swapchainImage.views = static_cast<VkImageView *>(malloc(sizeof(VkImageView) * vk.swapchainImage.imageCount));
    for(uint32_t i = 0; i < vk.swapchainImage.imageCount; i++){
        VkHelper::createImageView(vk.deviceInfo.device, vk.swapchainImage.images[i],
                                  VK_IMAGE_VIEW_TYPE_2D, vk.swapchainParam.format.format, VK_IMAGE_ASPECT_COLOR_BIT, &vk.swapchainImage.views[i]);
    }
    initImages();
    VkHelper::createRenderPass(vk.deviceInfo.device, vk.swapchainParam, &vk.renderPass);
    vk.framebufferCount = vk.swapchainImage.imageCount;
    vk.framebuffers = static_cast<VkFramebuffer *>(malloc(sizeof(VkFramebuffer) * vk.framebufferCount));
    VkHelper::createFramebuffer(vk.deviceInfo.device, vk.renderPass, vk.swapchainParam.extent.width, vk.swapchainParam.extent.height,
                                vk.framebufferCount, vk.swapchainImage.views, vk.depthImageView, vk.colorImageView, vk.framebuffers);
}

void destroySwapchainAndResource(){
    for (size_t i = 0; i < vk.framebufferCount; i++) {
        vkDestroyFramebuffer(vk.deviceInfo.device, vk.framebuffers[i], VK_ALLOC);
    }
    for(uint32_t i = 0; i < vk.swapchainImage.imageCount; i++){
        vkDestroyImageView(vk.deviceInfo.device, vk.swapchainImage.views[i], VK_ALLOC);
    }
    free(vk.swapchainImage.views);
    free(vk.swapchainImage.images);
    vkDestroySwapchainKHR(vk.deviceInfo.device, vk.swapchain, VK_ALLOC);
}

void reCreateSwapchain(){
    vkDeviceWaitIdle(vk.deviceInfo.device);
    Print("Recreate swapchain start.");

    destroySwapchainAndResource();
    createSwapchainAndResource();
}

void init(){
    VkHelper::createInstance(true, &vk.instance, &vk.debugReport);
    createWindow();
    VkHelper::pickPhyDevAndCreateDev(vk.instance, vk.surface, &vk.deviceInfo, &vk.queueInfo);
    vk.swapchainParam.maxSampleCount = VK_SAMPLE_COUNT_1_BIT;//VkHelper::getUsableSampleCountBits(vk.physicalDevLimits.framebufferColorSampleCounts);
    Print("Use Sample Count is %d", vk.swapchainParam.maxSampleCount);
    vk.maxFrameCount = MAX_FRAMES_COUNT;
    VkHelper::createCommandPool(vk.deviceInfo.device, vk.queueInfo.workQueueIndex, &vk.cmdPool);
    createSwapchainAndResource();
    VkHelper::createPipelineCache(vk.deviceInfo.device, &vk.pipelineCache);
    VkHelper::createDescriptorSetLayout(vk.deviceInfo.device, &vk.descriptorSetLayout);
    VkHelper::createDescriptorPool(vk.deviceInfo.device, vk.maxFrameCount, &vk.descriptorPool);
    vk.descriptorSets = static_cast<VkDescriptorSet *>(malloc(sizeof(VkDescriptorSet) * vk.maxFrameCount));
    VkHelper::allocateDescriptorSets(vk.deviceInfo.device, vk.maxFrameCount, vk.descriptorPool, vk.descriptorSetLayout, vk.descriptorSets);
    VkHelper::createPipelineLayout(vk.deviceInfo.device, vk.descriptorSetLayout, &vk.pipelineLayout);

    std::string shaders("shaders/");
    auto vertexShaderCode = readFile(shaders + defaultShader + ".vert.spv");
    auto fragShaderCode = readFile(shaders + defaultShader + ".frag.spv");
    vk.vertexShaderModule = VkHelper::createShaderModule(vk.deviceInfo.device, vertexShaderCode);
    vk.fragShaderModule = VkHelper::createShaderModule(vk.deviceInfo.device, fragShaderCode);
    VkHelper::createPipeline(vk.deviceInfo.device, vk.pipelineCache, vk.pipelineLayout, vk.renderPass,
                             vk.vertexShaderModule, vk.fragShaderModule, vk.swapchainParam, &vk.graphicPipeline);
    vk.cmdBufferCount = vk.swapchainImage.imageCount;
    vk.cmdBuffers = static_cast<VkCommandBuffer *>(malloc(sizeof(VkCommandBuffer) * vk.cmdBufferCount));
    VkHelper::allocateCommandBuffers(vk.deviceInfo.device, vk.cmdPool, vk.cmdBufferCount, vk.cmdBuffers);

    //int geometry, include vertex, texcoord, normal...
    initGeometry();
    vk.imageSemaphores = static_cast<VkSemaphore *>(malloc(sizeof(VkSemaphore) * vk.maxFrameCount));
    vk.presentSemaphores = static_cast<VkSemaphore *>(malloc(sizeof(VkSemaphore) * vk.maxFrameCount));
    vk.fences = static_cast<VkFence *>(malloc(sizeof(VkFence) * vk.maxFrameCount));
    VkSemaphoreCreateInfo semaphoreCreateInfo = {
            .sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0
    };
    VkFenceCreateInfo fenceCreateInfo = {
            .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0
    };
    for(uint32_t i = 0; i < vk.maxFrameCount; i++){
        CALL_VK(vkCreateSemaphore(vk.deviceInfo.device, &semaphoreCreateInfo, VK_ALLOC, &vk.imageSemaphores[i]));
        CALL_VK(vkCreateSemaphore(vk.deviceInfo.device, &semaphoreCreateInfo, VK_ALLOC, &vk.presentSemaphores[i]));
        CALL_VK(vkCreateFence(vk.deviceInfo.device, &fenceCreateInfo, VK_ALLOC, &vk.fences[i]));
    }
}

void render(uint32_t currentFrame){

//        vkWaitForFences(vk.deviceInfo.device, 1, &vk.fences[currentFrame], VK_TRUE, std::numeric_limits<uint64_t>::max());
//        vkResetFences(vk.deviceInfo.device, 1, &vk.fences[currentFrame]);
    uint32_t imageIndex = 0;
    VkResult rt = vkAcquireNextImageKHR(vk.deviceInfo.device, vk.swapchain, std::numeric_limits<uint64_t>::max(),
                                  vk.imageSemaphores[currentFrame], VK_NULL_HANDLE, &imageIndex);
    if(rt == VK_ERROR_OUT_OF_DATE_KHR){
        reCreateSwapchain();
        return;
    } else if (rt != VK_SUCCESS && rt != VK_SUBOPTIMAL_KHR) {
        throw std::runtime_error("failed to acquire swap chain image!");
    }

    //update
    updateUniformBuffer(currentFrame);
    updateDescriptorSets(currentFrame);
    VkClearValue defaultClearValues[] = {
            {
                .color = {0.1f, 0.1f, 0.2f, 1.0f}
            },
            {
                .depthStencil = {1, 0}
            }
    };
    for(uint32_t i = 0; i < vk.cmdBufferCount; i++){
        VkHelper::beginCommandBuffer(vk.cmdBuffers[i], false);
        VkHelper::beginRenderPass(vk.cmdBuffers[i], vk.renderPass, vk.framebuffers[i], ARRAY_SIZE(defaultClearValues), defaultClearValues, vk.swapchainParam.extent);
        VkHelper::bindRenderResource(vk.cmdBuffers[i], vk.pipelineLayout, vk.graphicPipeline, vk.descriptorSets[currentFrame], geometry);
        VkHelper::geometryDraw(vk.cmdBuffers[i], vk.swapchainParam, geometry);
        VkHelper::endRenderPass(vk.cmdBuffers[i]);
        VkHelper::endCommandBuffer(vk.cmdBuffers[i], vk.deviceInfo.device, vk.cmdPool, vk.queueInfo.queue, false);
    }

    VkPipelineStageFlags stages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
    VkSubmitInfo submitInfo = {
            .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO,
            .pNext = nullptr,
            .waitSemaphoreCount = 1,
            .pWaitSemaphores = &vk.imageSemaphores[currentFrame],
            .pWaitDstStageMask = stages,
            .commandBufferCount = 1,
            .pCommandBuffers = &vk.cmdBuffers[imageIndex],
            .signalSemaphoreCount = 1,
            .pSignalSemaphores = &vk.presentSemaphores[currentFrame],
    };
    CALL_VK(vkQueueSubmit(vk.queueInfo.queue, 1, &submitInfo, VK_NULL_HANDLE));//vk.fences[currentFrame]

    VkPresentInfoKHR presentInfo = {
            .sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
            .pNext = nullptr,
            .waitSemaphoreCount = 1,
            .pWaitSemaphores = &vk.presentSemaphores[currentFrame],
            .swapchainCount = 1,
            .pSwapchains = &vk.swapchain,
            .pImageIndices = &imageIndex,
            .pResults = nullptr
    };
    rt = vkQueuePresentKHR(vk.queueInfo.queue, &presentInfo);
    if(rt == VK_ERROR_OUT_OF_DATE_KHR){
        reCreateSwapchain();
        return;
    } else if (rt != VK_SUCCESS && rt != VK_SUBOPTIMAL_KHR) {
        throw std::runtime_error("failed to acquire swap chain image!");
    }

    CALL_VK(vkQueueWaitIdle(vk.queueInfo.queue));
}

void renderLoop(){
    uint32_t currentFrame = 0;
    while(!glfwWindowShouldClose(window)){
        glfwPollEvents();

        render(currentFrame);

        glfwSwapBuffers(window);
        currentFrame = (currentFrame + 1) % vk.maxFrameCount;

        std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }
}

void destroy(){
    Print("%s", __FUNCTION__);
    vkDeviceWaitIdle(vk.deviceInfo.device);

    vkDestroyImageView(vk.deviceInfo.device, vk.colorImageView, VK_ALLOC);
    vkDestroyImage(vk.deviceInfo.device, vk.colorImage, VK_ALLOC);
    vkFreeMemory(vk.deviceInfo.device, vk.colorImageMemory, VK_ALLOC);
    texture2.destroy(vk.deviceInfo.device);
    texture1.destroy(vk.deviceInfo.device);
    geometry.destroy(vk.deviceInfo.device);
    vkFreeDescriptorSets(vk.deviceInfo.device, vk.descriptorPool, vk.maxFrameCount, vk.descriptorSets);
    free(vk.descriptorSets);
    vkDestroyDescriptorPool(vk.deviceInfo.device, vk.descriptorPool, VK_ALLOC);
    vkDestroyDescriptorSetLayout(vk.deviceInfo.device, vk.descriptorSetLayout, VK_ALLOC);
    for(uint32_t i = 0; i < vk.maxFrameCount; i++){
        vkDestroyFence(vk.deviceInfo.device, vk.fences[i], VK_ALLOC);
        vkDestroySemaphore(vk.deviceInfo.device, vk.presentSemaphores[i], VK_ALLOC);
        vkDestroySemaphore(vk.deviceInfo.device, vk.imageSemaphores[i], VK_ALLOC);
    }
    free(vk.fences);
    free(vk.presentSemaphores);
    free(vk.imageSemaphores);
    for(uint32_t i = 0; i < vk.framebufferCount; i++){
        vkDestroyFramebuffer(vk.deviceInfo.device, vk.framebuffers[i], VK_ALLOC);
    }
    vkFreeCommandBuffers(vk.deviceInfo.device, vk.cmdPool, vk.cmdBufferCount, vk.cmdBuffers);
    free(vk.cmdBuffers);
    free(vk.framebuffers);
    vkFreeMemory(vk.deviceInfo.device, vk.depthImageMemory, VK_ALLOC);
    vkDestroyImage(vk.deviceInfo.device, vk.depthImage, VK_ALLOC);
    vkDestroyImageView(vk.deviceInfo.device, vk.depthImageView, VK_ALLOC);
    vkDestroySampler(vk.deviceInfo.device, vk.depthImageSampler, VK_ALLOC);
    vkDestroyCommandPool(vk.deviceInfo.device, vk.cmdPool, VK_ALLOC);
    vkDestroyPipeline(vk.deviceInfo.device, vk.graphicPipeline, VK_ALLOC);
    vkDestroyShaderModule(vk.deviceInfo.device, vk.fragShaderModule, VK_ALLOC);
    vkDestroyShaderModule(vk.deviceInfo.device, vk.vertexShaderModule, VK_ALLOC);
    vkDestroyPipelineLayout(vk.deviceInfo.device, vk.pipelineLayout, VK_ALLOC);
    vkDestroyPipelineCache(vk.deviceInfo.device, vk.pipelineCache, VK_ALLOC);
    vkDestroyRenderPass(vk.deviceInfo.device, vk.renderPass, VK_ALLOC);
    destroySwapchainAndResource();
    vkDestroyDevice(vk.deviceInfo.device, VK_ALLOC);
    vkDestroySurfaceKHR(vk.instance, vk.surface, VK_ALLOC);
    vkDestroyInstance(vk.instance, VK_ALLOC);
    glfwDestroyWindow(window);
    glfwTerminate();
}

static void onFramebufferSizeCallback(GLFWwindow *glfWwindow, int width, int height){
    if(width != vk.swapchainParam.extent.width || height != vk.swapchainParam.extent.height){
        reCreateSwapchain();
    }
}

void createWindow(){
    if(!glfwInit()){
        throw std::runtime_error("Can not init glfw window.");
    }
    glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
    window = glfwCreateWindow(WINDOW_WIDTH, WINDOW_HEIGHT, APP_NAME, nullptr, nullptr);
    if(!window){
        throw std::runtime_error("Can not create glfw window");
    }

    Print("GLFW Vulkan Support: %d", glfwVulkanSupported());
    glfwMakeContextCurrent(window);

    if(glfwCreateWindowSurface(vk.instance, window, VK_ALLOC, &vk.surface) != VK_SUCCESS){
        throw std::runtime_error("Can not creat surface.");
    }

    Print("surface: %p", vk.surface);

    glfwSetFramebufferSizeCallback(window, &onFramebufferSizeCallback);
}

void initImages(){
    vk.swapchainParam.depthFormat = VkHelper::findDepthSupportFormat(vk.deviceInfo.physicalDev);
    VkHelper::createImage(vk.deviceInfo.physicalDevMemoProps, vk.deviceInfo.device, vk.swapchainParam.extent.width, vk.swapchainParam.extent.height, 1,
                          VK_IMAGE_TYPE_2D, vk.swapchainParam.depthFormat, vk.swapchainParam.maxSampleCount, VK_IMAGE_TILING_OPTIMAL, VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT,
                          &vk.depthImage, &vk.depthImageMemory);
    VkHelper::createImageView(vk.deviceInfo.device, vk.depthImage, VK_IMAGE_VIEW_TYPE_2D, vk.swapchainParam.depthFormat,
                              VK_IMAGE_ASPECT_DEPTH_BIT, &vk.depthImageView);
    VkHelper::transitionImageLayout(vk.deviceInfo.device, vk.cmdPool, vk.queueInfo.queue, vk.depthImage, vk.swapchainParam.depthFormat, VK_IMAGE_ASPECT_DEPTH_BIT,
                                    VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL);

    VkFormat colorFormat = vk.swapchainParam.format.format;
    VkHelper::createImage(vk.deviceInfo.physicalDevMemoProps, vk.deviceInfo.device, vk.swapchainParam.extent.width, vk.swapchainParam.extent.height, 1,
                          VK_IMAGE_TYPE_2D, colorFormat, vk.swapchainParam.maxSampleCount, VK_IMAGE_TILING_OPTIMAL,
                          VK_IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
                          &vk.colorImage, &vk.colorImageMemory);
    VkHelper::createImageView(vk.deviceInfo.device, vk.colorImage, VK_IMAGE_VIEW_TYPE_2D, colorFormat,
                              VK_IMAGE_ASPECT_COLOR_BIT, &vk.colorImageView);
    VkHelper::transitionImageLayout(vk.deviceInfo.device, vk.cmdPool, vk.queueInfo.queue, vk.colorImage, vk.swapchainParam.depthFormat, VK_IMAGE_ASPECT_COLOR_BIT,
                                    VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
//    VkHelper::transitionImageLayout(vk.deviceInfo.device, vk.cmdPool, vk.queue, vk.colorImage, vk.swapchainParam.depthFormat, VK_IMAGE_ASPECT_COLOR_BIT,
//                                    VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
}

void initGeometry(){
//    geometry.vertices = {
//            {{0.5f, 0.5f, 0.0f},/*vertex*/ {1.0f, 0.0f},/*texcoord*/    {1.0f, 0.0f, 0.0f},/*normal*/},
//            {{0.0f, 0.5f, 0.0f},/*vertex*/ {1.0f, 1.0f},/*texcoord*/    {0.0f, 1.0f, 0.0f},/*normal*/},
//            {{0.0f, 0.0f, 0.0f},/*vertex*/ {0.0f, 1.0f},/*texcoord*/    {0.0f, 0.0f, 1.0f},/*normal*/},
//            {{0.5f, 0.0f, 0.0f},/*vertex*/ {0.0f, 0.0f},/*texcoord*/    {1.0f, 1.0f, 1.0f},/*normal*/},
//    };
//    geometry.indices = {
//            0, 1, 2, 0, 2, 3
//    };
    if(!getCube(&geometry, 0.5f)){
        throw std::runtime_error("build geometry fail.");
    }
    VkHelper::initGeometryBuffers(vk.deviceInfo.physicalDevMemoProps, vk.deviceInfo.device, vk.cmdPool, vk.queueInfo.queue, vk.maxFrameCount, geometry);

    texture1.load(vk.deviceInfo.physicalDevMemoProps, vk.deviceInfo.physicalDevLimits, vk.deviceInfo.device, vk.cmdPool, vk.queueInfo.queue, "texture/container.jpg");
    texture2.load(vk.deviceInfo.physicalDevMemoProps, vk.deviceInfo.physicalDevLimits, vk.deviceInfo.device, vk.cmdPool, vk.queueInfo.queue, "texture/awesomeface.png");
}

void updateUniformBuffer(uint32_t currentFrame){
    auto time = (float) glfwGetTime();
    Print("%.3f", time);
    ubo.model = glm::rotate(glm::mat4(1.0f), time * 2.0f, glm::vec3(1.0f, 1.0f, 0.0f));
    ubo.model = glm::translate(glm::mat4(1.0f), glm::vec3(cos(time), sin(time), 0));
    ubo.view = glm::lookAt(glm::vec3(2.0f, 2.0f, 2.0f), glm::vec3(0.0f, 0.0f, 0.0f), glm::vec3(0.0f, 0.0f, 1.0f));
    ubo.proj = glm::perspective(glm::radians(45.0f),
                                (float) vk.swapchainParam.extent.width / (float) vk.swapchainParam.extent.height, 0.1f, 10.0f);

    memcpy(geometry.uniformBuffersMapped[currentFrame], &ubo, sizeof(ubo));
}

void updateDescriptorSets(uint32_t currentframe){
    VkDescriptorBufferInfo descBufferInfo = {
            .buffer = geometry.uniformBuffers[currentframe],
            .offset = 0,
            .range = sizeof(UniformObject)
    };

    VkDescriptorImageInfo imageInfo1[] = {
            {
                .sampler = texture1.imageSampler,
                .imageView = texture1.imageView,
                .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
            }
    };

    VkDescriptorImageInfo imageInfo2[] = {
            {
                    .sampler = texture2.imageSampler,
                    .imageView = texture2.imageView,
                    .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
            }
    };

    VkWriteDescriptorSet writeDescSets[] = {
            {
                .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                .pNext = nullptr,
                .dstSet = vk.descriptorSets[currentframe],
                .dstBinding = 0,
                .dstArrayElement = 0,
                .descriptorCount = 1,
                .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                .pBufferInfo = &descBufferInfo
            },
            {
                .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                .pNext = nullptr,
                .dstSet = vk.descriptorSets[currentframe],
                .dstBinding = 1,
                .dstArrayElement = 0,
                .descriptorCount = ARRAY_SIZE(imageInfo1),
                .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                .pImageInfo = imageInfo1
            },
            {
                .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
                .pNext = nullptr,
                .dstSet = vk.descriptorSets[currentframe],
                .dstBinding = 2,
                .dstArrayElement = 0,
                .descriptorCount = ARRAY_SIZE(imageInfo2),
                .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                .pImageInfo = imageInfo2
            }
    };
    vkUpdateDescriptorSets(vk.deviceInfo.device, ARRAY_SIZE(writeDescSets), writeDescSets, 0, nullptr);
}